import sys
import cv2
import threading
import time
import os
import numpy as np
from datetime import datetime
import mss
import pyautogui
from PIL import Image
import io
import base64
import pyaudio
import wave
import ffmpeg
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                            QHBoxLayout, QPushButton, QLabel, QMessageBox, 
                            QListWidget, QFrame)
from PyQt6.QtCore import Qt, QTimer, pyqtSignal, QSize
from PyQt6.QtGui import QIcon, QPixmap, QImage

os.environ["QT_AUTO_SCREEN_SCALE_FACTOR"] = "1"

class VideoRecorder:
    def __init__(self):
        self.is_recording = False
        self.cap = None
        self.out = None
        self.current_file = None
        self._lock = threading.Lock()
        self.recording_type = "camera"
        self.preview_frame = None
        self.camera_error = False
        self.audio_recorder = AudioRecorder()
        
    def start_recording(self, recording_type="camera"):
        with self._lock:
            try:
                self.is_recording = True
                self.recording_type = recording_type
                
                if recording_type == "camera":
                    if self.cap is None:
                        self.cap = cv2.VideoCapture(0)
                        # 添加摄像头初始化检查
                        if not self.cap.isOpened():
                            self.camera_error = True
                            raise Exception("无法打开摄像头，请检查摄像头是否正确连接")
                        
                        # 尝试读取一帧来验证摄像头是否正常工作
                        ret, _ = self.cap.read()
                        if not ret:
                            self.camera_error = True
                            raise Exception("无法从摄像头读取画面")
                            
                    frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                    frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                else:  # screen
                    screen_size = pyautogui.size()
                    frame_width, frame_height = screen_size.width, screen_size.height
                
                fps = 20.0
                
                if not os.path.exists("recordings"):
                    os.makedirs("recordings")
                    
                filename = f"recordings/{'screen' if recording_type == 'screen' else 'camera'}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
                self.current_file = filename
                
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                self.out = cv2.VideoWriter(filename, fourcc, fps, (frame_width, frame_height))
                self.audio_recorder.start_recording()

            except Exception as e:
                self.is_recording = False
                self.audio_recorder.is_recording = False
                if self.cap is not None:
                    self.cap.release()
                    self.cap = None
                raise e

    def capture_screen(self):
        with mss.mss() as sct:  # 在函数内部创建 mss 实例
            monitor = sct.monitors[1]  # 主显示器
            screenshot = sct.grab(monitor)
            img = Image.frombytes('RGB', screenshot.size, screenshot.rgb)
            frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            return frame

    def get_preview_frame(self):
        if not self.is_recording:
            return None
            
        try:
            if self.recording_type == "camera" and self.cap:
                if self.camera_error:
                    return None
                    
                ret, frame = self.cap.read()
                if not ret:
                    self.camera_error = True
                    print("无法获取摄像头画面")
                    return None
                    
                # 调整预览尺寸
                frame = cv2.resize(frame, (640, 480))
                return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            elif self.recording_type == "screen":
                frame = self.capture_screen()
                # 调整预览尺寸
                frame = cv2.resize(frame, (640, 480))
                return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        except Exception as e:
            print(f"获取预览画面时出错: {e}")
            return None
        return None
        
    def stop_recording(self):
        with self._lock:
            self.is_recording = False
            self.audio_recorder.stop_recording()
            if self.out is not None:
                self.out.release()
                self.out = None
            if self.cap is not None:
                self.cap.release()
                self.cap = None
            
            # 合并音视频文件
            if self.current_file and self.audio_recorder.current_file:
                try:
                    output_file = self.current_file.replace('.mp4', '_with_audio.mp4')
                    
                    # 使用 ffmpeg 合并音视频
                    input_video = ffmpeg.input(self.current_file)
                    input_audio = ffmpeg.input(self.audio_recorder.current_file)
                    
                    ffmpeg.output(
                        input_video,
                        input_audio,
                        output_file,
                        vcodec='copy',
                        acodec='aac'
                    ).overwrite_output().run(capture_stdout=True, capture_stderr=True)
                    
                    # 清理临时文件
                    os.remove(self.current_file)
                    os.remove(self.audio_recorder.current_file)
                    
                    # 更新当前文件路径
                    self.current_file = output_file
                    
                except Exception as e:
                    print(f"合并音视频时出错: {e}")
    
    def __del__(self):
        try:
            if self.out is not None:
                self.out.release()
            if self.cap is not None:
                self.cap.release()
        except Exception as e:
            print(f"清理资源时出错: {e}")

class AudioRecorder:
    def __init__(self):
        self.is_recording = False
        self.frames = []
        self._lock = threading.Lock()
        self.current_file = None
        
    def start_recording(self):
        self.is_recording = True
        self.frames = []
        if not os.path.exists("recordings"):
            os.makedirs("recordings")
        self.current_file = f"recordings/audio_{datetime.now().strftime('%Y%m%d_%H%M%S')}.wav"
        
    def record_audio(self):
        CHUNK = 1024
        FORMAT = pyaudio.paInt16
        CHANNELS = 2
        RATE = 44100
        
        p = pyaudio.PyAudio()
        stream = p.open(format=FORMAT,
                       channels=CHANNELS,
                       rate=RATE,
                       input=True,
                       frames_per_buffer=CHUNK)
                       
        while self.is_recording:
            try:
                data = stream.read(CHUNK)
                with self._lock:
                    self.frames.append(data)
            except Exception as e:
                print(f"录制音频时出错: {e}")
                break
                
        stream.stop_stream()
        stream.close()
        p.terminate()
        
        # 保存音频文件
        if self.frames:
            wf = wave.open(self.current_file, 'wb')
            wf.setnchannels(CHANNELS)
            wf.setsampwidth(p.get_sample_size(FORMAT))
            wf.setframerate(RATE)
            wf.writeframes(b''.join(self.frames))
            wf.close()
            
    def stop_recording(self):
        self.is_recording = False

class MainWindow(QMainWindow):
    update_preview_signal = pyqtSignal(QImage)
    
    def __init__(self):
        super().__init__()
        self.recorder = VideoRecorder()
        self.recording_thread = None
        self.preview_thread = None
        self.recording_time = 0
        self._is_stopping = False  # 添加停止标志
        self.initUI()
        
    def initUI(self):
        self.setWindowTitle('视频录制软件')
        self.setFixedSize(800, 800)
        
        # 创建主窗口部件和布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # 状态标签
        self.status_label = QLabel('就绪')
        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.status_label.setStyleSheet('font-size: 20px;')
        layout.addWidget(self.status_label)
        
        # 计时器标签
        self.timer_label = QLabel('00:00')
        self.timer_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.timer_label.setStyleSheet('font-size: 30px;')
        layout.addWidget(self.timer_label)
        
        # 预览窗口
        self.preview_label = QLabel("预览窗口")
        self.preview_label.setFixedSize(640, 480)
        self.preview_label.setFrameStyle(QFrame.Shape.Box)
        self.preview_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.preview_label.setStyleSheet("QLabel { background-color: #f0f0f0; }")
        layout.addWidget(self.preview_label)
        
        # 按钮布局
        button_layout = QHBoxLayout()
        
        self.camera_btn = QPushButton('摄像头录制')
        self.camera_btn.setIcon(QIcon.fromTheme('camera-video'))
        self.camera_btn.clicked.connect(lambda: self.toggle_recording('camera'))
        button_layout.addWidget(self.camera_btn)
        
        self.screen_btn = QPushButton('屏幕录制')
        self.screen_btn.setIcon(QIcon.fromTheme('display'))
        self.screen_btn.clicked.connect(lambda: self.toggle_recording('screen'))
        button_layout.addWidget(self.screen_btn)
        
        self.stop_btn = QPushButton('停止录制')
        self.stop_btn.setIcon(QIcon.fromTheme('media-playback-stop'))
        self.stop_btn.clicked.connect(self.stop_recording)
        self.stop_btn.hide()
        button_layout.addWidget(self.stop_btn)
        
        layout.addLayout(button_layout)
        
        # 分割线
        line = QFrame()
        line.setFrameShape(QFrame.Shape.HLine)
        line.setFrameShadow(QFrame.Shadow.Sunken)
        layout.addWidget(line)
        
        # 文件列表
        list_label = QLabel('录制文件列表')
        list_label.setStyleSheet('font-size: 16px; font-weight: bold;')
        layout.addWidget(list_label)
        
        self.files_list = QListWidget()
        layout.addWidget(self.files_list)
        
        self.update_files_list()
        
        # 设置预览信号连接
        self.update_preview_signal.connect(self.update_preview_label)
        
        # 设置定时器更新UI
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_timer_label)
        
    def toggle_recording(self, recording_type):
        if not self.recorder.is_recording:
            try:
                self.recorder.start_recording(recording_type)
                self.status_label.setText("录制中...")
                self.preview_label.show()
                
                # 更新按钮状态
                self.camera_btn.hide()
                self.screen_btn.hide()
                self.stop_btn.show()
                
                # 启动线程
                self.recording_thread = threading.Thread(target=self.record_video, daemon=True)
                self.preview_thread = threading.Thread(target=self.update_preview, daemon=True)
                
                self.recording_thread.start()
                self.preview_thread.start()
                
                # 启动计时器
                self.recording_time = 0
                self.timer.start(1000)  # 每秒更新一次
                
                # 启动音频录制线程
                self.audio_thread = threading.Thread(target=self.recorder.audio_recorder.record_audio, daemon=True)
                self.audio_thread.start()
                
                # 强制更新UI
                QApplication.processEvents()
                
            except Exception as e:
                self.status_label.setText(str(e) if str(e) else "录制启动失败")
                self.recorder.is_recording = False
                QApplication.processEvents()
                
    def stop_recording(self):
        try:
            # 设置停止标志
            self._is_stopping = True
            
            # 先停止录制
            self.recorder.stop_recording()
            
            # 等待线程结束
            if self.recording_thread and self.recording_thread.is_alive():
                self.recording_thread.join(timeout=1.0)
            if self.preview_thread and self.preview_thread.is_alive():
                self.preview_thread.join(timeout=1.0)
            
            # 停止计时器
            self.timer.stop()
            
            # 等待音频线程结束
            if hasattr(self, 'audio_thread') and self.audio_thread.is_alive():
                self.audio_thread.join(timeout=1.0)
            
            # 更新UI
            self.status_label.setText("视频已保存")
            self.timer_label.setText("00:00")
            self.preview_label.clear()
            self.preview_label.setText("预览窗口")
            self.recording_time = 0
            
            # 恢复按钮状态
            self.stop_btn.hide()
            self.camera_btn.show()
            self.screen_btn.show()
            
            # 更新文件列表
            self.update_files_list()
            
            # 重置停止标志
            self._is_stopping = False
            
            # 强制更新UI
            QApplication.processEvents()
            
        except Exception as e:
            print(f"停止录制时出错: {e}")
            self.status_label.setText("录制停止失败")
            # 确保按钮状态正确
            self.stop_btn.hide()
            self.camera_btn.show()
            self.screen_btn.show()
            self._is_stopping = False
            QApplication.processEvents()
            
    def update_preview(self):
        while self.recorder.is_recording and not self._is_stopping:
            try:
                frame = self.recorder.get_preview_frame()
                if frame is not None:
                    height, width, channel = frame.shape
                    bytes_per_line = 3 * width
                    q_image = QImage(frame.data, width, height, bytes_per_line, QImage.Format.Format_RGB888)
                    self.update_preview_signal.emit(q_image)
                time.sleep(0.033)
            except Exception as e:
                print(f"预览更新出错: {e}")
                time.sleep(0.1)
        print("预览线程结束")
                
    def update_preview_label(self, image):
        pixmap = QPixmap.fromImage(image)
        self.preview_label.setPixmap(pixmap.scaled(
            self.preview_label.size(),
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation
        ))
        
    def update_timer_label(self):
        self.recording_time += 1
        minutes = self.recording_time // 60
        seconds = self.recording_time % 60
        self.timer_label.setText(f"{minutes:02d}:{seconds:02d}")
        
    def record_video(self):
        while self.recorder.is_recording and not self._is_stopping:
            try:
                with self.recorder._lock:
                    if self.recorder.is_recording and not self._is_stopping:
                        if self.recorder.recording_type == "camera":
                            if self.recorder.cap and self.recorder.out:
                                ret, frame = self.recorder.cap.read()
                                if ret:
                                    self.recorder.out.write(frame)
                        else:  # screen
                            if self.recorder.out:
                                frame = self.recorder.capture_screen()
                                self.recorder.out.write(frame)
                time.sleep(0.001)
            except Exception as e:
                print(f"录制视频时出错: {e}")
                self.recorder.is_recording = False
                break
        print("录制线程结束")
                
    def update_files_list(self):
        self.files_list.clear()
        if os.path.exists("recordings"):
            for file in sorted(os.listdir("recordings"), reverse=True):
                if file.endswith(".mp4"):
                    self.files_list.addItem(file)
                    
    def closeEvent(self, event):
        reply = QMessageBox.question(
            self, '确认退出',
            '确定要退出视频录制软件吗？',
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No
        )

        if reply == QMessageBox.StandardButton.Yes:
            try:
                if self.recorder.is_recording:
                    # 停止录制
                    self.recorder.stop_recording()
                
                # 强制设置录制状态为False
                self.recorder.is_recording = False
                
                # 关闭视频流
                if self.recorder.out is not None:
                    self.recorder.out.release()
                    self.recorder.out = None
                if self.recorder.cap is not None:
                    self.recorder.cap.release()
                    self.recorder.cap = None
                
                # 等待线程结束
                if self.recording_thread and self.recording_thread.is_alive():
                    self.recording_thread.join(timeout=0.5)
                if self.preview_thread and self.preview_thread.is_alive():
                    self.preview_thread.join(timeout=0.5)
                
                # 停止计时器
                self.timer.stop()
                
                # 等待音频线程结束
                if hasattr(self, 'audio_thread') and self.audio_thread.is_alive():
                    self.audio_thread.join(timeout=0.5)
                
                event.accept()
            except Exception as e:
                print(f"关闭窗口时出错: {e}")
                event.accept()
        else:
            event.ignore()

def main():
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())

if __name__ == "__main__":
    main()